Two scalar-prefetch refs replace an assumption that only one request is
in flight:

- kv_new_starts_ref: per-seq base offset into the new-KV buffer. Without
  it every sequence reads its current K/V from offset 0, which is only
  correct when the launch carries a single request.
- kv_write_seq_mask_ref: which sequences perform the cache write. This
  replaces the write_last_seq_only static flag, which could express "the
  last of one request's two fused seqs" but not "the tail seq of each of
  R requests".

write_last_seq_only is removed rather than kept alongside: it is a
special case of the mask, and leaving both would let callers set
contradictory values.

Validation rejects kv_new_starts together with pcp_chunk_size — the
chunk-size remap rewrites addresses on the assumption of one request, so
the two are mutually exclusive by construction.

Tests cover per-request KV writes landing on the right pages exactly
once, and per-request current-phase output against a reference. One test
is skipped by default: it asserts the kernel tolerates a zero-length
sequence, which it does not — it hangs the device rather than raising,
so running it needs PCP_RUN_HANG_REPRO=1 under an external timeout.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

Cache phase (in-kernel ring): one seq per REQUEST spanning its head+tail
run.  Q stays local under the ring, so each rank's buffer is its request
slots back to back and the description is just the current-phase arrays
with the head/tail duplication undone: cu_q_lens, kv_lens, kv_cache_lens
and the block-table rows at even indices, and half the seq count from
request_distribution.  No tiling, no output collective.  R == 1 keeps
the ring PR's whole-buffer seq verbatim, which is valid for any
current-phase cu_q_lens including the clipped-tail, per-rank one older
callers build.

The R > 1 construction relies on the ring's lock-step: remote copies and
semaphores are matched by loop position, so every rank must run the
identical (seq, tile, block, round) schedule.  request_distribution is
replicated, cu_q_lens is rank-invariant because the runner gives both
halves the full chunk length, and the block count comes from the
replicated global cache length.

Current phase: the all-gathered K/V arrives in rank order.  With one
request the kernel remaps addresses itself via pcp_chunk_size; that
remap assumes a single request, so several requests reorder into
request-major order in JAX via kv_token_order and pass kv_new_starts
instead.  Each request's tail seq performs the fused KV write via
kv_write_seq_mask.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

Lifts the one-request-per-step restriction. Each request is zigzag-
chunked independently, so rank r holds chunk r and chunk 2P-1-r of every
request, and request i's head/tail pair occupies a fixed-width slot in
every rank's region.

The pieces:

- Token permutation from natural order to rank order, and its inverse
  for the K/V reorder. Slots past a request's end stay -1 and are zeroed.
- Attention metadata over 2R fused seqs, request i at seqs 2i and 2i+1.
  Both halves get the full chunk length, so cu_q_lens is rank-invariant
  and only q_pos_offsets varies by rank.
- logits_indices pointing at each request's last real token, wherever
  the zigzag put it.
- Block-table rows duplicated per fused seq, since the KV write is done
  by the tail seq and page addressing is per-seq.

Buffer sizing changes with it: attention seq slots double under PCP, and
the token bucket ladder gains 2*P*max_num_seqs of headroom. Rounding
each chunk up independently means P*S can exceed the scheduler's token
budget, and the scheduler admits by sum(n_i), so without that headroom a
legal batch has no bucket to land in.

A step with exactly ONE request keeps the single-request layout: chunk
t_pad / 2P, head and tail each filling half of every rank's slice, i.e.
exactly the layout upstream's runner produces. pcp_forward's R == 1
current phase (the kernel-side rank-order remap when the chunk is
page-aligned, to_token_order otherwise) derives the chunk from the padded
buffer width -- pcp_chunk_size is a static kernel argument -- so a lone
request packed at ceil(n / 2P) would have its current K/V reordered
against the wrong chunk whenever the token bucket is not exact. Multi-
request steps carry their own permutation (kv_token_order) and do not
depend on the buffer width. The extra tail rows past the request's last
token are padding rows like any other: they attend only real keys,
nothing reads their output, and the KV write is driven by kv_lens -
kv_cache_lens, not by the row count.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

Adds num_reqs to the warm-up ladder and builds well-formed dummy PCP
metadata for it.

The dummy metadata has to be a batch the runtime could actually produce.
An all-zero query_start_loc satisfies every shape check but describes
zero-length sequences, which hang the kernel on device during warm-up —
the failure looks like a deadlock at startup with no traceback.

The ladder is capped at {1, max_num_seqs} rather than every power of two,
and combinations the runtime cannot reach are skipped: cache_pages only
matters to the gather-KV path, which multi-request does not take, and a
batch of R requests cannot have fewer than 2*P*R tokens. Without these
the graph count roughly quadrupled and warm-up dominated startup.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

…{0, max}

Both attention phases now iterate the live seq count from
request_distribution, so the slots that pad the request count up to the
static num_reqs bucket never reach the kernel: drop their 1-token chunks
from the layout and the max(1, ...) floor in pcp_token_layout.

With the ring, cache_pages only decides whether the cache phase is
elided (== 0); every nonzero rung compiled the same kernel.  Two rungs
halve the PCP precompile set.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

Kernel (tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py):
- fused current phase with several requests: per-request kv_new_starts,
  kv_write_seq_mask, every request's strided KV share written exactly once,
  ragged lengths, mixed cached/uncached;
- ring cache phase over three seqs in one launch with different cached
  lengths (one zero), row counts and page lists, P = 2 and 4, f32 and bf16,
  against per-seq full-cache reference (per-seq block counts, one
  cross-rank chain across seq boundaries, lock-step).

Interface (tests/layers/common/test_pcp_attention_interface.py):
- pcp_forward with R requests: ragged lengths, mixed cached/uncached, live
  count below the static num_reqs bucket, per-request KV cache write, a
  request spanning several ring query tiles, the batch shape from a 32k
  end-to-end run in the headroom token bucket, and ~1M cached tokens next to
  a short request;
- single-request regression cases unchanged.

All need >= P TPU devices; run through cdk.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>

Review follow-ups for vllm-project#3425.

Branching: `kv_new_starts` and `kv_token_order` are now required
fields of PCPMetadata and are built by the runner, the precompile
dummies, the interface tests and the pcp-vs-tp benchmark for every
request count. That removes the optional-argument plumbing in
pcp_forward (`*pcp_extra`, `extra_specs`/`extra_args`, the None guard)
and the `if multi else None` in the runner and precompile. The
single-request `to_token_order`/`inv_row` reshuffle is gone too: the
runner's `kv_token_order` for one request is exactly that permutation,
so the non-page-aligned path is one `jnp.take` for any request count.
The page-aligned single-request fast path (kernel-side remap via
`pcp_chunk_size`) and the single-request cache-phase span are kept as
they were; merging them into the general layout is the follow-up PR.

`cache_pages` becomes `has_cached_kv: bool`. After the in-kernel ring
the value only decided whether the cache phase runs, so the two-rung
page ladder, its bucket constant and the round-up helper are deleted;
precompile loops over {False, True}.

Removed the multi-request batch logger and its counter, which were
debugging aids. Comments cut to what the code does not say itself.

Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>